#![cfg_attr(debug_assertions, allow(dead_code, unused_imports, unused_variables))]
use proc_macro::TokenStream;
use protocol::{SymbolicTask, TaskSignature};
use quote::quote;
use syn::{parse_macro_input, ItemStruct, Ident, Lit, Meta, parse::Parse, Token, punctuated::Punctuated, parse::ParseStream, parse::Result, Expr, ExprLit, parse::ParseBuffer, ExprArray, LitStr};
mod parsing;
mod protocol;
fn define_app_state() -> proc_macro2::TokenStream {
    quote! {
        // #[derive(Debug, Copy, Clone)]
        // pub enum TaskStatus {
        //     Created,
        //     Running,
        //     Completed,
        //     Failed,
        //     Stopped,
        //     Closed,
        //     ResultReceived,
        //     FailedToGetResult
        // }
        struct AppState {
            task_progress: std::sync::Mutex<std::collections::HashMap<String, (f32, edp_protocol::TaskStatus)>>,
            task_results: std::sync::Mutex<std::collections::HashMap<String, Option<Value>>>,
            task_errors: std::sync::Mutex<std::collections::HashMap<String, Option<String>>>,
            logs: std::sync::Mutex<std::collections::HashMap<String, (Vec<String>, usize)>>
        }
        impl AppState {
            pub fn new() -> Self {
                Self {
                    task_progress: std::sync::Mutex::new(std::collections::HashMap::new()),
                    task_results: std::sync::Mutex::new(std::collections::HashMap::new()),
                    task_errors: std::sync::Mutex::new(std::collections::HashMap::new()),
                    logs: std::sync::Mutex::new(std::collections::HashMap::new()),
                }
            }
            pub fn create_task(&self, tid: String) {
                self.task_progress.lock().unwrap().insert(tid.clone(), (0.0, edp_protocol::TaskStatus::Running));
                self.task_results.lock().unwrap().insert(tid.clone(), None);
                self.logs.lock().unwrap().insert(tid, (vec![], 0));
            }
            pub fn update_task_progress(&self, tid: &str, progress: f32, status: edp_protocol::TaskStatus) {
                self.task_progress.lock().unwrap().entry(tid.to_string()).and_modify(|(p, s)| {
                    *p = progress;
                    *s = status;
                });
            }
            pub fn insert_task_result(&self, tid: String, result: Value) {
                self.task_results.lock().unwrap().insert(tid, Some(result));
            }
            pub fn insert_task_error(&self, tid: String, error: String) {
                self.task_errors.lock().unwrap().insert(tid, Some(error));
            }
            pub fn log(&self, tid: &str, log: String) {
                self.logs.lock().unwrap().entry(tid.to_string()).and_modify(|(logs, _)| {
                    logs.push(log);
                });
            }
        }
    }
}

struct ServiceAttr {
    name: String,
    desc: Option<String>,
    category: Option<String>,
    nominal_tasks: Option<Vec<String>>,
    task_sig: Option<Vec<(TaskSignature, bool)>>,
    // symbolic_tasks: Option<Vec<SymbolicTask>>
}

impl Parse for ServiceAttr {
    fn parse(input: &ParseBuffer) -> Result<Self> {
        // 使用 Punctuated 来解析逗号分隔的列表
        let content: Punctuated<Meta, Token![,]> =
            input.parse_terminated(Meta::parse, Token![,])?;
        let mut name = None;
        let mut desc = None;
        let mut category = None;
        let mut nominal_tasks = None;
        let mut task_sig = None;
        for meta in content {
            match meta {
                Meta::NameValue(nv) => {
                    let ident = nv.path.get_ident().ok_or_else(|| {
                        input.error("Expected identifier for attribute key")
                    })?;
                    let key = ident.to_string();
                    if key == "name" {
                        if let syn::Expr::Lit(expr_lit) = nv.value {
                            if let syn::Lit::Str(lit_str) = expr_lit.lit {
                                name = Some(lit_str.value());
                            } else {
                                return Err(input.error("`name` must be a string literal"));
                            }
                        } else {
                            return Err(input.error("`name` must be a string literal"));
                        }
                    } else if key == "desc" {
                        if let syn::Expr::Lit(expr_lit) = nv.value {
                            if let syn::Lit::Str(lit_str) = expr_lit.lit {
                                desc = Some(lit_str.value());
                            } else {
                                return Err(input.error("`desc` must be a string literal"));
                            }
                        } else {
                            return Err(input.error("`desc` must be a string literal"));
                        }
                    } else if key == "category" {
                        if let syn::Expr::Lit(expr_lit) = nv.value {
                            if let syn::Lit::Str(lit_str) = expr_lit.lit {
                                category = Some(lit_str.value());
                            } else {
                                return Err(input.error("`category` must be a string literal"));
                            }
                        } else {
                            return Err(input.error("`category` must be a string literal"));
                        }
                    } else if key == "tasks" {
                        if let syn::Expr::Array(expr_array) = nv.value {
                            let mut nominals = Vec::new();
                            let mut sigs = Vec::new();
                            for elem_expr in expr_array.elems {
                                if let syn::Expr::Lit(elem_expr_lit) = elem_expr {
                                    if let syn::Lit::Str(s) = elem_expr_lit.lit {
                                        let (s, ins) = match s.value().trim() {
                                            s if s.starts_with("#instant ") => {
                                                (s[9..].to_owned(), true)
                                            }
                                            s => (s.to_owned(), false)
                                        };
                                        let symbolic_task = SymbolicTask::parse(s.as_str());
                                        nominals.push(s);
                                        sigs.push((symbolic_task.sig(), ins));
                                    } else {
                                        return Err(input.error("`tasks` array elements must be string literals"));
                                    }
                                } else {
                                    return Err(input.error("`tasks` array elements must be string literals"));
                                }
                            }
                            nominal_tasks = Some(nominals);
                            task_sig = Some(sigs);
                        } else {
                            return Err(input.error("`tasks` must be an array literal"));
                        }
                    } else {
                        return Err(input.error(format!("Unknown attribute key: {}", key)));
                    }
                }
                _ => {
                    return Err(input.error("Unsupported attribute format. Expected `key = value`"));
                }
            }
        }
        let name = name.ok_or_else(|| input.error("Missing required attribute `name`"))?;
        Ok(ServiceAttr {
            name,
            desc,
            category,
            nominal_tasks,
            task_sig
        })
    }
}

#[proc_macro_attribute]
pub fn edp_service(attr: TokenStream, item: TokenStream) -> TokenStream {
    let input = parse_macro_input!(item as ItemStruct);
    let name = &input.ident;
    let service_attr = parse_macro_input!(attr as ServiceAttr);
    let name_fn = {
        let val = service_attr.name;
        quote! { fn name() -> &'static str { #val } }
    };
    let desc_fn = if let Some(val) = service_attr.desc {
        quote! { fn desc() -> &'static str { #val } }
    } else { quote!{} };
    let category_fn = if let Some(val) = service_attr.category {
        quote! { fn category() -> Option<&'static str> { Some(#val) } }
    } else { quote!{ fn category() -> Option<&'static str> { None } } };
    let tasks_fn = if let Some(val) = service_attr.nominal_tasks {
        let tasks = val.iter().map(|s| quote! { #s });
        quote! { 
            fn tasks() -> Vec<&'static str> { 
                vec![#(#tasks),*]
            } 
        }
    } else { quote!{ fn tasks() -> Vec<&'static str> { vec![] } } };
    let new_fn = quote! {
        fn new(task_id: String, state: actix_web::web::Data<std::sync::Arc<AppState>>) -> Self {
            Self {
                task_id,
                state,
            }
        }
    };
    let instant_fn = quote! {
        fn _instant() -> Self {
            Self {
                task_id: "".into(),
                state: actix_web::web::Data::new(std::sync::Arc::new(AppState::new())),
            }
        }
    };
    let update_fn = quote! {
        pub fn update_progress(&self, progress: f32, status: edp_protocol::TaskStatus) {
            self.state.update_task_progress(&self.task_id, progress, status);
        }
    };
    let log_fn = quote! {
        pub fn log(&self, log: String) {
            self.state.log(&self.task_id, log);
        }
    };
    let mut long_tasks = Vec::new();
    let mut instant_tasks = Vec::new();
    if let Some(val) = service_attr.task_sig {
        for each in val {
            if each.1 {
                instant_tasks.push(each.0)
            } else {
                long_tasks.push(each.0)
            }
        }
    }
    // instant
    let instant_routes = instant_tasks.iter().map(|t|{
        let task_name = syn::Ident::new(t.name.as_str(), proc_macro2::Span::call_site());
        let fn_name: Ident = syn::Ident::new(&format!("start_{}", t.name), proc_macro2::Span::call_site());
        let route = format!("/task/{}", t.name);
        let input_type_str = t.input_type.as_str();
        let get_input_type = quote! {
            &edp_protocol::TypeSchema::parse(#input_type_str)
        };
        let get_result = quote! {
            #name::_instant().#task_name(data)
        };
        println!("expand instant task: {:?}", t);
        quote! {
            #[actix_web::post(#route)]
            async fn #fn_name(req: actix_web::HttpRequest, state: actix_web::web::Data<std::sync::Arc<AppState>>, body: actix_web::web::Json<serde_json::Value>) -> actix_web::HttpResponse {
                check_session_key!(&req);
                println!("in this func: {}", #route);
                let data = body.into_inner();
                if !edp_protocol::qualify(&data, #get_input_type) {
                    return actix_web::HttpResponse::BadRequest().json(serde_json::json!({ "error": "input data qualification failed" }));
                }
                let res = #get_result;
                match res {
                    Ok(res) => {
                        actix_web::HttpResponse::Ok().json(serde_json::json!({
                            "task_id": "",
                            "progress_pipe_addr": "",
                            "result_addr": "",
                            "log_pipe_addr": "",
                            "instant_result": res,
                            "status": "ResultReceived"
                        }))
                    }
                    Err(e) => {
                        actix_web::HttpResponse::Ok().json(serde_json::json!({
                            "task_id": "",
                            "progress_pipe_addr": "",
                            "result_addr": "",
                            "log_pipe_addr": "",
                            "instant_result": null,
                            "status": "Failed",
                            "error_message": e
                        }))
                    }
                }
            }
        }
    });
    let instant_routes: Vec<proc_macro2::TokenStream> = instant_routes.collect();
    let service_instant = instant_tasks.iter().map(|t|{
        let fn_name = syn::Ident::new(&format!("start_{}", t.name), proc_macro2::Span::call_site());
        quote! {
            .service(#fn_name)
        }
    });
    let service_instant: Vec<proc_macro2::TokenStream> = service_instant.collect();
    // long
    let long_start_fn = long_tasks.iter().map(|t|{
        let fn_name = syn::Ident::new(&format!("_{}", t.name), proc_macro2::Span::call_site());
        let inner_fn_name = syn::Ident::new(&format!("{}", t.name), proc_macro2::Span::call_site());
        quote! {
            fn #fn_name (&mut self, inputs: serde_json::Value) {
                self.state.create_task(self.task_id.clone());
                let result = self.#inner_fn_name(inputs);
                println!("task result: {:?}", result);
                match result {
                    Ok(result) => {
                        self.update_progress(100.0, edp_protocol::TaskStatus::Completed);
                        self.state.insert_task_result(self.task_id.clone(), result);
                    }
                    Err(e) => {
                        self.update_progress(0.0, edp_protocol::TaskStatus::Failed);
                        self.state.insert_task_error(self.task_id.clone(), e);
                    }
                }
            }
        }
    });
    let long_start_fn: Vec<proc_macro2::TokenStream> = long_start_fn.collect();
    let long_routes = long_tasks.iter().map(|t|{
        let route = format!("/task/{}", t.name);
        let fn_name = syn::Ident::new(&format!("start_{}", t.name), proc_macro2::Span::call_site());
        let task_fn_name = syn::Ident::new(&format!("_{}", t.name), proc_macro2::Span::call_site());
        let input_type_str = t.input_type.as_str();
        let get_input_type = quote! {
            &edp_protocol::TypeSchema::parse(#input_type_str)
        };
        let progress_pipe_addr_lead = format!("/task/{}/progress/", t.name);
        let result_addr_lead = format!("/task/{}/result/", t.name);
        let log_pipe_addr_lead = format!("/task/{}/log/", t.name);
        let start_route = quote! {
            #[actix_web::post(#route)]
            async fn #fn_name (req: actix_web::HttpRequest, state: actix_web::web::Data<std::sync::Arc<AppState>>, body: actix_web::web::Json<serde_json::Value>) -> actix_web::HttpResponse {
                check_session_key!(&req);
                let data = body.into_inner();
                if !edp_protocol::qualify(&data, #get_input_type) {
                    return actix_web::HttpResponse::BadRequest().json(serde_json::json!({ "error": "input data qualification failed" }));
                }
                let task_id = uuid::Uuid::new_v4().to_string();
                let mut task = #name::new(task_id.clone(), state.clone());
                std::thread::spawn(move || {
                    task.#task_fn_name(data);
                });
                actix_web::HttpResponse::Ok().json(serde_json::json!({
                    "task_id": task_id,
                    "progress_pipe_addr": #progress_pipe_addr_lead.to_string() + &task_id,
                    "result_addr": #result_addr_lead.to_string() + &task_id,
                    "log_pipe_addr": #log_pipe_addr_lead.to_string() + &task_id,
                }))
            }
        };
        let progress_route_val = format!("{}{{task_id}}", progress_pipe_addr_lead);
        let fn_name = syn::Ident::new(&format!("get_{}_progress", t.name), proc_macro2::Span::call_site());
        let progress_route = quote! {
            #[actix_web::get(#progress_route_val)]
            async fn #fn_name (task_id: actix_web::web::Path<String>, state: actix_web::web::Data<std::sync::Arc<AppState>>) -> actix_web::HttpResponse {
                if let Some(task) = state.task_progress.lock().unwrap().get(&task_id.into_inner()) {
                    return actix_web::HttpResponse::Ok().json(serde_json::json!({
                        "percentage": task.0,
                        "status": format!("{:?}", task.1)
                    }).to_string())
                } else {
                    return actix_web::HttpResponse::NotFound().json(serde_json::json!({ "error": "Task not found" }))
                }
            }
        };
        let result_route_val = format!("{}{{task_id}}", result_addr_lead);
        // println!("{}", result_route_val);
        let fn_name = syn::Ident::new(&format!("get_{}_result", t.name), proc_macro2::Span::call_site());
        let result_route = quote! {
            #[actix_web::get(#result_route_val)]
            async fn #fn_name (task_id: actix_web::web::Path<String>, state: actix_web::web::Data<std::sync::Arc<AppState>>) -> actix_web::HttpResponse {
                if let Some(result) = state.task_results.lock().unwrap().get(&task_id.into_inner()) {
                    if let Some(result) = result {
                        return actix_web::HttpResponse::Ok().json(result)
                    }
                    return actix_web::HttpResponse::BadRequest().json(serde_json::json!({ "error": "Task not finished" }))
                } else {
                    return actix_web::HttpResponse::NotFound().json(serde_json::json!({ "error": "Task not found" }))
                }
            }
        };
        let log_route_val = format!("{}{{task_id}}", log_pipe_addr_lead);
        let fn_name = syn::Ident::new(&format!("get_{}_log", t.name), proc_macro2::Span::call_site());
        let log_route = quote! {
            #[actix_web::get(#log_route_val)]
            async fn #fn_name (task_id: actix_web::web::Path<String>, state: actix_web::web::Data<std::sync::Arc<AppState>>) -> actix_web::HttpResponse {
                let mut logs_guard = state.logs.lock().unwrap();
                let tid = task_id.into_inner();
                if !logs_guard.contains_key(&tid) {
                    return actix_web::HttpResponse::NotFound().json(serde_json::json!({ "error": "Task not found" }))
                }
                let logs = logs_guard.entry(tid).or_insert((vec![], 0));
                let newlogs = logs.0[logs.1..].to_vec();
                logs.1 = logs.0.len();
                return actix_web::HttpResponse::Ok().json(serde_json::json!(newlogs))
            }
        };
        println!("expand long task: {:?}", t);
        vec![start_route, progress_route, result_route, log_route]
    });
    let long_routes: Vec<Vec<proc_macro2::TokenStream>> = long_routes.collect();
    let service_long = long_tasks.iter().map(|t|{
        let start_fn_name = syn::Ident::new(&format!("start_{}", t.name), proc_macro2::Span::call_site());
        let progress_fn_name = syn::Ident::new(&format!("get_{}_progress", t.name), proc_macro2::Span::call_site());
        let result_fn_name = syn::Ident::new(&format!("get_{}_result", t.name), proc_macro2::Span::call_site());
        let log_fn_name = syn::Ident::new(&format!("get_{}_log", t.name), proc_macro2::Span::call_site());
        quote!{
            .service(#start_fn_name)
            .service(#progress_fn_name)
            .service(#result_fn_name)
            .service(#log_fn_name)
        }
    });
    let service_long: Vec<proc_macro2::TokenStream> = service_long.collect();
    let index_route = quote! {
        #[actix_web::get("/")]
        async fn index() -> String {
            #name::name().to_string()
        }
    };
    let meta_route = quote! {
        #[actix_web::get("/meta")]
        async fn meta() -> String {
            json!({
                "name": #name::name(),
                "desc": #name::desc(),
                "category": #name::category(),
            }).to_string()
        }
    };
    let tasks_route = quote! {
        #[actix_web::get("/tasks")]
        async fn tasks() -> String {
            json!(#name::tasks()).to_string()
        }
    };
    let workflows_route = quote! {
        #[actix_web::get("/workflows")]
        async fn workflows() -> String {
            json!(#name::workflows()).to_string()
        }
    };
    let mount_route = quote! {
        #[actix_web::get("/mount")]
        async fn mount() -> String {
            json!({
                "session_key": uuid::Uuid::new_v4().to_string(),
                "healthy_info": "Healthy",
            }).to_string()
        }
    };
    let service_fn = quote! {
        pub async fn service(port: u16) -> std::io::Result<()> {
            let state = std::sync::Arc::new(AppState::new());
            actix_web::HttpServer::new(move || {
                actix_web::App::new()
                    .app_data(actix_web::web::Data::new(state.clone()))
                    .service(index)
                    .service(meta)
                    .service(tasks)
                    .service(workflows)
                    .service(mount)
                    #(#service_instant)*
                    #(#service_long)*
            })
            .bind(format!("127.0.0.1:{}", port))?
            .run()
            .await
        }
    };
    let check_session_key_macro = quote! {
        macro_rules! check_session_key {
            ($req:expr) => {
                // if $req.headers().get("session_key").is_none() {
                //     return HttpResponse::BadRequest().json(json!({ "error": "session key not found" }));
                // }
            };
        }
    };
    let app_state = define_app_state();
    let expanded = quote! {
        #app_state

        struct #name {
            pub task_id: String,
            pub state: actix_web::web::Data<std::sync::Arc<AppState>>
        }
        
        impl #name {
            #name_fn
            #desc_fn
            #category_fn
            #tasks_fn
            fn workflows() -> Vec<&'static str> { vec![] }
            #new_fn
            #instant_fn
            #update_fn
            #log_fn
            #(#long_start_fn)*
        }
        #check_session_key_macro
        #index_route
        #meta_route
        #tasks_route
        #workflows_route
        #mount_route
        #(#instant_routes)*
        #(#(#long_routes)*)*
        #service_fn
    };

    TokenStream::from(expanded)
} 

#[proc_macro_attribute]
pub fn instant(_attr: TokenStream, item: TokenStream) -> TokenStream {
    item
}